package cn.lgwen.candy.connector.elasticsearch2.sink;

import cn.lgwen.candy.runtime.operator.SinkOperator;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.common.transport.InetSocketTransportAddress;

import java.util.LinkedList;
import java.util.List;

/**
 * 2019/12/30
 * aven.wu
 * danxieai258@163.com
 */
public abstract class Elasticsearch2Sink extends SinkOperator<IndexRequest> {

    private TransportClient client;

    private int bulkSize = 1000;

    private List<IndexRequest> indexRequests = new LinkedList<>();

    private final Object lock = new Object();

    public Elasticsearch2Sink(List<InetSocketTransportAddress> addressList) {
        client = TransportClient.builder().build()
                .addTransportAddresses(addressList.toArray(new InetSocketTransportAddress[]{}));
    }

    public void setBulkSize(int bulkSize) {
        this.bulkSize = bulkSize;
    }

    @Override
    public void invoke(IndexRequest input) {
        synchronized (lock) {
            if (indexRequests.size() < bulkSize) {
                indexRequests.add(input);
            } else {
                BulkRequest bulkRequest = new BulkRequest();
                indexRequests.forEach(bulkRequest::add);
                client.bulk(bulkRequest);
                indexRequests.clear();
            }
        }
    }
}
